# Description:
#   This package contains ScaNN partitioners.

load("//scann/oss_wrappers:build_defs.bzl", "scann_py_proto_library")

package(
    default_visibility = ["//visibility:public"],
    licenses = ["notice"],
)

# Binaries
# ==============================================================================

# Libraries
# ==============================================================================

cc_library(
    name = "anisotropic",
    srcs = ["anisotropic.cc"],
    hdrs = ["anisotropic.h"],
    tags = ["local"],
    deps = [
        "//scann/oss_wrappers:tf_dependency",
        "//scann/utils:common",
        "//scann/utils/linear_algebra:eigen_utils",
        "//scann/utils/linear_algebra:types",
    ],
)

cc_library(
    name = "partitioner_base",
    srcs = ["partitioner_base.cc"],
    hdrs = ["partitioner_base.h"],
    tags = ["local"],
    deps = [
        ":partitioner_cc_proto",
        "//scann/data_format:datapoint",
        "//scann/data_format:dataset",
        "//scann/oss_wrappers:scann_aligned_malloc",
        "//scann/oss_wrappers:scann_down_cast",
        "//scann/oss_wrappers:scann_status",
        "//scann/oss_wrappers:scann_threadpool",
        "//scann/oss_wrappers:tf_dependency",
        "//scann/utils:types",
        "//scann/utils:util_functions",
        "//scann/utils:zip_sort",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/flags:flag",
        "@com_google_absl//absl/log",
    ],
)

cc_library(
    name = "orthogonality_amplification_utils",
    hdrs = ["orthogonality_amplification_utils.h"],
    tags = ["local"],
    deps = [
        "//scann/data_format:datapoint",
        "//scann/utils:types",
    ],
)

cc_library(
    name = "kmeans_tree_like_partitioner",
    hdrs = ["kmeans_tree_like_partitioner.h"],
    tags = ["local"],
    deps = [
        ":partitioner_base",
        "//scann/data_format:datapoint",
        "//scann/distance_measures",
        "//scann/oss_wrappers:scann_threadpool",
        "//scann/trees/kmeans_tree",
        "//scann/utils:types",
        "@com_google_absl//absl/status",
    ],
)

cc_library(
    name = "kmeans_tree_partitioner",
    srcs = [
        "kmeans_tree_partitioner.cc",
        "kmeans_tree_partitioner_helper.cc",
    ],
    hdrs = ["kmeans_tree_partitioner.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_like_partitioner",
        ":kmeans_tree_partitioner_cc_proto",
        ":orthogonality_amplification_utils",
        ":partitioner_base",
        ":partitioner_cc_proto",
        "//scann/base:search_parameters",
        "//scann/base:single_machine_base",
        "//scann/brute_force",
        "//scann/data_format:dataset",
        "//scann/distance_measures/many_to_many",
        "//scann/hashes/asymmetric_hashing2:indexing",
        "//scann/hashes/asymmetric_hashing2:querying",
        "//scann/hashes/asymmetric_hashing2:searcher",
        "//scann/hashes/asymmetric_hashing2:training",
        "//scann/hashes/asymmetric_hashing2:training_options",
        "//scann/oss_wrappers:scann_aligned_malloc",
        "//scann/oss_wrappers:scann_down_cast",
        "//scann/oss_wrappers:scann_status",
        "//scann/oss_wrappers:scann_threadpool",
        "//scann/oss_wrappers:tf_dependency",
        "//scann/proto:distance_measure_cc_proto",
        "//scann/proto:hash_cc_proto",
        "//scann/proto:partitioning_cc_proto",
        "//scann/trees/kmeans_tree",
        "//scann/utils:common",
        "//scann/utils:datapoint_utils",
        "//scann/utils:fast_top_neighbors",
        "//scann/utils:parallel_for",
        "//scann/utils:reordering_helper",
        "//scann/utils:types",
        "//scann/utils:zip_sort",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/flags:flag",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "tree_brute_force_second_level_wrapper",
    srcs = ["tree_brute_force_second_level_wrapper.cc"],
    hdrs = ["tree_brute_force_second_level_wrapper.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_like_partitioner",
        "//scann/brute_force",
        "//scann/oss_wrappers:scann_status",
        "//scann/proto:partitioning_cc_proto",
        "//scann/tree_x_hybrid:tree_x_hybrid_smmd",
        "//scann/trees/kmeans_tree",
        "//scann/utils:datapoint_utils",
        "//scann/utils:types",
    ],
)

cc_library(
    name = "kmeans_tree_partitioner_utils",
    hdrs = ["kmeans_tree_partitioner_utils.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_partitioner",
        ":partitioner_base",
        "//scann/distance_measures",
        "//scann/oss_wrappers:scann_status",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "projecting_decorator",
    srcs = ["projecting_decorator.cc"],
    hdrs = ["projecting_decorator.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_like_partitioner",
        ":partitioner_base",
        "//scann/oss_wrappers:scann_status",
        "//scann/oss_wrappers:scann_threadpool",
        "//scann/projection:projection_base",
        "//scann/utils:common",
        "//scann/utils:datapoint_utils",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
    ],
)

cc_library(
    name = "partitioner_factory_base",
    srcs = ["partitioner_factory_base.cc"],
    hdrs = ["partitioner_factory_base.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_partitioner_utils",
        ":partitioner_base",
        ":projecting_decorator",
        "//scann/data_format:dataset",
        "//scann/distance_measures",
        "//scann/oss_wrappers:scann_status",
        "//scann/projection:projection_base",
        "//scann/projection:projection_factory",
        "//scann/proto:distance_measure_cc_proto",
        "//scann/proto:partitioning_cc_proto",
        "//scann/utils:types",
        "//scann/utils:util_functions",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/time",
    ],
)

cc_library(
    name = "partitioner_factory",
    srcs = ["partitioner_factory.cc"],
    hdrs = ["partitioner_factory.h"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_partitioner_cc_proto",
        ":kmeans_tree_partitioner_utils",
        ":partitioner_base",
        ":partitioner_cc_proto",
        ":partitioner_factory_base",
        ":projecting_decorator",
        "//scann/data_format:dataset",
        "//scann/distance_measures",
        "//scann/oss_wrappers:scann_random",
        "//scann/oss_wrappers:scann_status",
        "//scann/projection:projection_base",
        "//scann/projection:projection_factory",
        "//scann/proto:distance_measure_cc_proto",
        "//scann/proto:partitioning_cc_proto",
        "//scann/trees/kmeans_tree",
        "//scann/utils:common",
        "//scann/utils:types",
        "//scann/utils:util_functions",
        "//scann/utils:weak_ptr_cache",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/time",
    ],
)

proto_library(
    name = "partitioner_proto",
    srcs = ["partitioner.proto"],
    tags = ["local"],
    deps = [
        ":kmeans_tree_partitioner_proto",
        ":linear_projection_tree_proto",
        "//scann/data_format:features_proto",
    ],
)

cc_proto_library(
    name = "partitioner_cc_proto",
    tags = ["local"],
    deps = [":partitioner_proto"],
)

scann_py_proto_library(
    name = "partitioner_py_pb2",
    srcs = ["partitioner.proto"],
    proto_deps = [":partitioner_proto"],
    py_proto_deps = [
        ":kmeans_tree_partitioner_py_pb2",
        ":linear_projection_tree_py_pb2",
    ],
    tags = ["local"],
)

proto_library(
    name = "kmeans_tree_partitioner_proto",
    srcs = ["kmeans_tree_partitioner.proto"],
    tags = ["local"],
    deps = [
        "//scann/proto:partitioning_proto",
        "//scann/trees/kmeans_tree:kmeans_tree_proto",
    ],
)

cc_proto_library(
    name = "kmeans_tree_partitioner_cc_proto",
    tags = ["local"],
    deps = [":kmeans_tree_partitioner_proto"],
)

scann_py_proto_library(
    name = "kmeans_tree_partitioner_py_pb2",
    srcs = ["kmeans_tree_partitioner.proto"],
    proto_deps = [":kmeans_tree_partitioner_proto"],
    py_proto_deps = [
        "//scann/proto:partitioning_py_pb2",
        "//scann/trees/kmeans_tree:kmeans_tree_py_pb2",
    ],
    tags = ["local"],
)

proto_library(
    name = "linear_projection_tree_proto",
    srcs = ["linear_projection_tree.proto"],
    tags = ["local"],
    deps = [
    ],
)

scann_py_proto_library(
    name = "linear_projection_tree_py_pb2",
    srcs = ["linear_projection_tree.proto"],
    proto_deps = [":linear_projection_tree_proto"],
    tags = ["local"],
)

# Tests
# =============================================================================
